import torch
import numpy
import math

def rotation_from_vector(approach_vector, closing_vector):
    #TODO: if element in approach_vector[:, 2] = 0 cause error !
    temp = -(approach_vector[:, 0] * closing_vector[:, 0] + approach_vector[:, 1] * closing_vector[:, 1]) / approach_vector[:, 2]
    closing_vector[:, 2] = temp
    closing_vector = torch.div(closing_vector.transpose(0, 1), torch.norm(closing_vector, dim=1)).transpose(0, 1)
    z_axis = torch.cross(approach_vector, closing_vector, dim = 1)
    R = torch.stack((approach_vector, closing_vector, z_axis), dim=-1)
    return R

def grasp_angle_to_vector(grasp_angle):
    device = grasp_angle.device
    x_ = torch.cos(grasp_angle/180*math.pi)
    y_ = torch.sin(grasp_angle/180*math.pi)
    z_ = torch.zeros(grasp_angle.shape[0]).to(device)
    closing_vector = torch.stack((x_, y_, z_), axis=0)

    return closing_vector

def refine_grasp_angle_to_vector(grasp_angle):
    device = grasp_angle.device
    x_ = torch.cos((grasp_angle+90)/180*math.pi)
    y_ = torch.sin((grasp_angle+90)/180*math.pi)
    z_ = torch.zeros(grasp_angle.shape[0]).to(device)
    closing_vector = torch.stack((x_, y_, z_), axis=0)

    return closing_vector

def angle_to_vector(azimuth_angle, elevation_angle):
    device = azimuth_angle.device
    x_ = torch.cos(azimuth_angle/180*math.pi)
    y_ = torch.sin(azimuth_angle/180*math.pi)
    z_ = -torch.tan((elevation_angle+15)/180 * math.pi) * torch.sqrt(x_**2 + y_**2)
    approach_vector_ = torch.stack((x_, y_, z_), axis=0)
    #print(torch.norm(approach_vector_, dim=0).shape)
    #print(approach_vector_.shape)
    #exit()
    approach_vector = torch.div(approach_vector_, torch.norm(approach_vector_, dim=0))
    return approach_vector

def decode_bbox_target(pred_reg, loc_scope, loc_bin_size, azimuth_scope, azimuth_bin_size, elevation_scope, elevation_bin_size, width_scope, width_bin_size, grasp_angle_scope, grasp_angle_bin_size):
    """
    :param roi_box3d: (N, 7)
    :param pred_reg: (N, C)
    :param loc_scope:
    :param loc_bin_size:
    :param loc_y_scope:
    :param loc_bin_size:
    :param get_ry_fine:
    :return:
    """
    per_loc_bin_num = int(loc_scope / loc_bin_size) * 2

    # recover xz localization
    x_bin_l, x_bin_r = 0, per_loc_bin_num
    z_bin_l, z_bin_r = per_loc_bin_num, per_loc_bin_num * 2
    start_offset = z_bin_r

    x_bin = torch.argmax(pred_reg[:, x_bin_l: x_bin_r], dim=1)
    z_bin = torch.argmax(pred_reg[:, z_bin_l: z_bin_r], dim=1)

    #pos_x = x_bin.float() * loc_bin_size + loc_bin_size / 2 - loc_scope
    #pos_z = z_bin.float() * loc_bin_size + loc_bin_size / 2 - loc_scope
    pos_x = x_bin.float() * loc_bin_size - loc_scope
    pos_z = z_bin.float() * loc_bin_size - loc_scope

    x_res_l, x_res_r = per_loc_bin_num * 2, per_loc_bin_num * 3
    z_res_l, z_res_r = per_loc_bin_num * 3, per_loc_bin_num * 4
    start_offset = z_res_r

    x_res_norm = torch.gather(torch.sigmoid(pred_reg[:, x_res_l: x_res_r]), dim=1, index=x_bin.unsqueeze(dim=1)).squeeze(dim=1)
    z_res_norm = torch.gather(torch.sigmoid(pred_reg[:, z_res_l: z_res_r]), dim=1, index=z_bin.unsqueeze(dim=1)).squeeze(dim=1)
    x_res = x_res_norm * loc_bin_size
    z_res = z_res_norm * loc_bin_size

    pos_x += x_res
    pos_z += z_res

    # recover y localization
    y_bin_l, y_bin_r = start_offset, start_offset + per_loc_bin_num
    y_res_l, y_res_r = y_bin_r, y_bin_r + per_loc_bin_num
    start_offset = y_res_r

    y_bin = torch.argmax(pred_reg[:, y_bin_l: y_bin_r], dim=1)
    y_res_norm = torch.gather(torch.sigmoid(pred_reg[:, y_res_l: y_res_r]), dim=1, index=y_bin.unsqueeze(dim=1)).squeeze(dim=1)
    y_res = y_res_norm * loc_bin_size
    #pos_y = y_bin.float() * loc_bin_size + loc_bin_size / 2 - loc_scope + y_res
    pos_y = y_bin.float() * loc_bin_size - loc_scope + y_res
    pos = torch.stack((pos_x, pos_y, pos_z), axis=1)

    per_azimuth_bin_num = int(azimuth_scope / azimuth_bin_size)
    per_elevation_bin_num = int(elevation_scope / elevation_bin_size)
    azimuth_bin_l, azimuth_bin_r = start_offset, start_offset + per_azimuth_bin_num
    azimuth_res_l, azimuth_res_r = azimuth_bin_r, azimuth_bin_r + per_azimuth_bin_num
    start_offset = azimuth_res_r
    azimuth_bin = torch.argmax(pred_reg[:, azimuth_bin_l: azimuth_bin_r], dim=1)
    azimuth_res_norm = torch.gather(torch.sigmoid(pred_reg[:, azimuth_res_l: azimuth_res_r]), dim=1, index=azimuth_bin.unsqueeze(dim=1)).squeeze(dim=1)
    azimuth_res = azimuth_res_norm * azimuth_bin_size
    #azimuth_angle = azimuth_bin.float() * azimuth_bin_size + azimuth_bin_size / 2 + azimuth_res
    azimuth_angle = azimuth_bin.float() * azimuth_bin_size + azimuth_res

    elevation_bin_l, elevation_bin_r = start_offset, start_offset + per_elevation_bin_num
    elevation_res_l, elevation_res_r = elevation_bin_r, elevation_bin_r + per_elevation_bin_num
    start_offset = elevation_res_r
    elevation_bin = torch.argmax(pred_reg[:, elevation_bin_l: elevation_bin_r], dim=1)
    elevation_res_norm = torch.gather(torch.sigmoid(pred_reg[:, elevation_res_l: elevation_res_r]), dim=1, index=elevation_bin.unsqueeze(dim=1)).squeeze(dim=1)
    elevation_res = elevation_res_norm * elevation_bin_size
    #elevation_angle = elevation_bin.float() * elevation_bin_size + elevation_bin_size / 2 + elevation_res
    elevation_angle = elevation_bin.float() * elevation_bin_size + elevation_res

    approach_vector = angle_to_vector(azimuth_angle, elevation_angle).transpose(0, 1) #vector B*N, 3
    #print(approach_vector[0])

    per_width_bin_num = int(width_scope / width_bin_size)
    width_bin_l, width_bin_r = start_offset, start_offset + per_width_bin_num
    width_res_l, width_res_r = width_bin_r, width_bin_r + per_width_bin_num
    start_offset = width_res_r
    width_bin = torch.argmax(pred_reg[:, width_bin_l: width_bin_r], dim=1)
    width_res_norm = torch.gather(torch.sigmoid(pred_reg[:, width_res_l: width_res_r]), dim=1, index=width_bin.unsqueeze(dim=1)).squeeze(dim=1)
    width_res = width_res_norm * width_bin_size
    #width = width_bin.float() * width_bin_size + width_bin_size / 2 + width_res
    width = width_bin.float() * width_bin_size + width_res

    per_grasp_angle_bin_num = int(grasp_angle_scope / grasp_angle_bin_size)
    grasp_angle_bin_l, grasp_angle_bin_r = start_offset, start_offset + per_grasp_angle_bin_num
    grasp_angle_res_l, grasp_angle_res_r = grasp_angle_bin_r, grasp_angle_bin_r + per_grasp_angle_bin_num
    start_offset = grasp_angle_res_r
    grasp_angle_bin = torch.argmax(pred_reg[:, grasp_angle_bin_l: grasp_angle_bin_r], dim=1)
    grasp_angle_res_norm = torch.gather(torch.sigmoid(pred_reg[:, grasp_angle_res_l: grasp_angle_res_r]), dim=1, index=grasp_angle_bin.unsqueeze(dim=1)).squeeze(dim=1)
    grasp_angle_res = grasp_angle_res_norm * grasp_angle_bin_size
    #grasp_angle = grasp_angle_bin.float() * grasp_angle_bin_size + grasp_angle_bin_size / 2 + grasp_angle_res
    grasp_angle = grasp_angle_bin.float() * grasp_angle_bin_size + grasp_angle_res
    #print(pos_x[0])
    #print(pos_y[0])
    #print(pos_z[0])
    #print(azimuth_angle[0])
    #print(elevation_angle[0])
    #print(grasp_angle[0])
    closing_vector = grasp_angle_to_vector(grasp_angle).transpose(0, 1)
    R = rotation_from_vector(approach_vector, closing_vector)
    #print(R[0])
    return pos, width, R

#    def canonical_transform_batch(pts_input, roi_boxes3d, gt_boxes3d):
#        """
#        :param pts_input: (N, npoints, 3 + C)
#        :param roi_boxes3d: (N, 7)
#        :param gt_boxes3d: (N, 7)
#        :return:
#        """
#        #roi_ry = roi_boxes3d[:, 6] % (2 * np.pi)  # 0 ~ 2pi
#        roi_center = roi_boxes3d[:, 0:3]
#        # shift to center
#        pts_input[:, :, [0, 1, 2]] = pts_input[:, :, [0, 1, 2]] - roi_center.reshape(-1, 1, 3)
#        gt_boxes3d_ct = np.copy(gt_boxes3d)
#        gt_boxes3d_ct[:, 0:3] = gt_boxes3d_ct[:, 0:3] - roi_center
#        # rotate to the direction of head
#        gt_boxes3d_ct = kitti_utils.rotate_pc_along_y_torch(torch.from_numpy(gt_boxes3d_ct.reshape(-1, 1, 7)),
#                                                            torch.from_numpy(roi_ry)).numpy().reshape(-1, 7)
#        gt_boxes3d_ct[:, 6] = gt_boxes3d_ct[:, 6] - roi_ry
#        pts_input = kitti_utils.rotate_pc_along_y_torch(torch.from_numpy(pts_input), torch.from_numpy(roi_ry)).numpy()
#
#        return pts_input, gt_boxes3d_ct
